#include "SystemServiceCapture.h"
#include "Utils.h"
#include <iostream>
#include <psapi.h>

SystemServiceCapture::SystemServiceCapture() :
    m_initialized(false),
    m_hToken(NULL),
    m_originalDesktop(NULL) {
}

SystemServiceCapture::~SystemServiceCapture() {
    Cleanup();
}

bool SystemServiceCapture::Initialize() {
    Utils::DebugLog("SystemServiceCapture: Initializing system service capture");

    // Get current process token
    if (!OpenProcessToken(GetCurrentProcess(), TOKEN_ALL_ACCESS, &m_hToken)) {
        m_lastError = "Failed to open process token: " + Utils::GetLastWin32Error();
        return false;
    }

    // Enable all required privileges
    if (!EnableAllPrivileges()) {
        Utils::DebugLog("SystemServiceCapture: Warning - Could not enable all privileges");
    }

    // Save original desktop
    m_originalDesktop = GetThreadDesktop(GetCurrentThreadId());

    m_initialized = true;
    Utils::DebugLog("SystemServiceCapture: Initialized successfully");
    return true;
}

void SystemServiceCapture::Cleanup() {
    if (m_originalDesktop && GetThreadDesktop(GetCurrentThreadId()) != m_originalDesktop) {
        SetThreadDesktop(m_originalDesktop);
    }

    if (m_hToken) {
        CloseHandle(m_hToken);
        m_hToken = NULL;
    }

    m_initialized = false;
}

bool SystemServiceCapture::EnableAllPrivileges() {
    const LPCSTR privileges[] = {
        SE_DEBUG_NAME,
        SE_TCB_NAME,
        SE_ASSIGNPRIMARYTOKEN_NAME,
        SE_LOCK_MEMORY_NAME,
        SE_INCREASE_QUOTA_NAME,
        SE_MACHINE_ACCOUNT_NAME,
        SE_SYSTEM_PROFILE_NAME,
        SE_PROF_SINGLE_PROCESS_NAME,
        SE_CREATE_PAGEFILE_NAME,
        SE_BACKUP_NAME,
        SE_RESTORE_NAME,
        SE_SHUTDOWN_NAME,
        SE_TAKE_OWNERSHIP_NAME,
        SE_LOAD_DRIVER_NAME,
        SE_SYSTEM_ENVIRONMENT_NAME,
        SE_CHANGE_NOTIFY_NAME,
        SE_UNDOCK_NAME,
        SE_MANAGE_VOLUME_NAME,
        SE_IMPERSONATE_NAME,
        SE_CREATE_GLOBAL_NAME,
        SE_ENABLE_DELEGATION_NAME,
        SE_TRUSTED_CREDMAN_ACCESS_NAME
    };

    TOKEN_PRIVILEGES tp;
    tp.PrivilegeCount = 1;
    tp.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;

    bool allSucceeded = true;

    for (int i = 0; i < sizeof(privileges) / sizeof(privileges[0]); i++) {
        if (LookupPrivilegeValueA(NULL, privileges[i], &tp.Privileges[0].Luid)) {
            if (!AdjustTokenPrivileges(m_hToken, FALSE, &tp, 0, NULL, NULL)) {
                Utils::DebugLog("SystemServiceCapture: Failed to enable privilege: " + std::string(privileges[i]));
                allSucceeded = false;
            }
        }
    }

    return allSucceeded;
}

bool SystemServiceCapture::CaptureActiveSession(HBITMAP& hBitmap, int& width, int& height) {
    if (!m_initialized) {
        m_lastError = "SystemServiceCapture not initialized";
        return false;
    }

    DWORD activeSessionId = WTSGetActiveConsoleSessionId();
    Utils::DebugLog("SystemServiceCapture: Active console session ID: " + std::to_string(activeSessionId));

    return CaptureSpecificSession(activeSessionId, hBitmap, width, height);
}

bool SystemServiceCapture::CaptureSpecificSession(DWORD sessionId, HBITMAP& hBitmap, int& width, int& height) {
    if (!m_initialized) {
        m_lastError = "SystemServiceCapture not initialized";
        return false;
    }

    Utils::DebugLog("SystemServiceCapture: Attempting to capture session " + std::to_string(sessionId));

    // Check if session is locked
    if (IsSessionLocked(sessionId)) {
        Utils::DebugLog("SystemServiceCapture: Session is locked - this is exactly what we want to capture!");
    }

    return CaptureSessionDesktop(sessionId, hBitmap, width, height);
}

bool SystemServiceCapture::IsSessionLocked(DWORD sessionId) {
    BOOL locked = FALSE;
    DWORD bytesReturned;

    if (WTSQuerySessionInformationA(WTS_CURRENT_SERVER_HANDLE, sessionId,
                                   WTSSessionInfoEx, NULL, &bytesReturned)) {
        // Session exists, try to determine if locked
        LPWSTR pSessionInfo = NULL;
        if (WTSQuerySessionInformationW(WTS_CURRENT_SERVER_HANDLE, sessionId,
                                       WTSSessionInfo, (LPWSTR*)&pSessionInfo, &bytesReturned)) {
            WTSFreeMemory(pSessionInfo);
        }
    }

    return locked;
}

bool SystemServiceCapture::CaptureSessionDesktop(DWORD sessionId, HBITMAP& hBitmap, int& width, int& height) {
    Utils::DebugLog("SystemServiceCapture: Capturing desktop for session " + std::to_string(sessionId));

    // Method 1: Try to impersonate the session user and capture
    HANDLE hUserToken = NULL;
    if (WTSQueryUserToken(sessionId, &hUserToken)) {
        Utils::DebugLog("SystemServiceCapture: Got user token for session");

        // Impersonate the user
        if (ImpersonateLoggedOnUser(hUserToken)) {
            Utils::DebugLog("SystemServiceCapture: Successfully impersonated session user");

            // Try to switch to session desktop
            if (SwitchToSessionDesktop(sessionId)) {
                bool result = CaptureCurrentDesktopContent(hBitmap, width, height);

                // Restore original desktop
                if (m_originalDesktop) {
                    SetThreadDesktop(m_originalDesktop);
                }

                RevertToSelf();
                CloseHandle(hUserToken);

                if (result) {
                    Utils::DebugLog("SystemServiceCapture: Session desktop capture succeeded with impersonation");
                    return true;
                }
            }

            RevertToSelf();
        }
        CloseHandle(hUserToken);
    }

    // Method 2: Direct desktop access without impersonation
    Utils::DebugLog("SystemServiceCapture: Trying direct desktop access");

    if (SwitchToSessionDesktop(sessionId)) {
        bool result = CaptureCurrentDesktopContent(hBitmap, width, height);

        // Restore original desktop
        if (m_originalDesktop) {
            SetThreadDesktop(m_originalDesktop);
        }

        if (result) {
            Utils::DebugLog("SystemServiceCapture: Session desktop capture succeeded with direct access");
            return true;
        }
    }

    // Method 3: Try to access all possible desktops in the session
    Utils::DebugLog("SystemServiceCapture: Trying to enumerate session desktops");

    const char* desktopNames[] = { "Winlogon", "Default", "Screen-saver", "Disconnect" };

    for (int i = 0; i < 4; i++) {
        char stationName[256];
        sprintf_s(stationName, "WinSta%d", sessionId);

        HWINSTA hWinStation = OpenWindowStationA(stationName, FALSE, WINSTA_ALL_ACCESS);
        if (!hWinStation) {
            hWinStation = OpenWindowStationA("WinSta0", FALSE, WINSTA_ALL_ACCESS);
        }

        if (hWinStation) {
            HDESK hDesktop = OpenDesktopA(desktopNames[i], 0, FALSE,
                DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS | DESKTOP_CREATEWINDOW | DESKTOP_CREATEMENU);

            if (hDesktop) {
                Utils::DebugLog("SystemServiceCapture: Opened desktop: " + std::string(desktopNames[i]));

                HDESK hOriginal = GetThreadDesktop(GetCurrentThreadId());
                if (SetThreadDesktop(hDesktop)) {
                    bool result = CaptureCurrentDesktopContent(hBitmap, width, height);
                    SetThreadDesktop(hOriginal);

                    CloseDesktop(hDesktop);
                    CloseWindowStation(hWinStation);

                    if (result) {
                        Utils::DebugLog("SystemServiceCapture: Successfully captured from desktop: " + std::string(desktopNames[i]));
                        return true;
                    }
                }
                CloseDesktop(hDesktop);
            }
            CloseWindowStation(hWinStation);
        }
    }

    m_lastError = "All session desktop capture methods failed";
    return false;
}

bool SystemServiceCapture::SwitchToSessionDesktop(DWORD sessionId) {
    // Try to get the input desktop for the session
    HDESK hInputDesktop = OpenInputDesktop(0, FALSE,
        DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS | DESKTOP_CREATEWINDOW);

    if (hInputDesktop) {
        if (SetThreadDesktop(hInputDesktop)) {
            Utils::DebugLog("SystemServiceCapture: Switched to input desktop");
            CloseDesktop(hInputDesktop);
            return true;
        }
        CloseDesktop(hInputDesktop);
    }

    // Try session-specific window station
    char stationName[256];
    sprintf_s(stationName, "WinSta%d", sessionId);

    HWINSTA hWinStation = OpenWindowStationA(stationName, FALSE, WINSTA_ALL_ACCESS);
    if (!hWinStation) {
        hWinStation = OpenWindowStationA("WinSta0", FALSE, WINSTA_ALL_ACCESS);
    }

    if (!hWinStation) {
        Utils::DebugLog("SystemServiceCapture: Failed to open window station");
        return false;
    }

    // Set the window station for current thread
    HWINSTA hOriginalStation = GetProcessWindowStation();
    SetProcessWindowStation(hWinStation);

    // Try to open the default desktop
    HDESK hDesktop = OpenDesktopA("Default", 0, FALSE,
        DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS | DESKTOP_CREATEWINDOW);

    if (!hDesktop) {
        // Try Winlogon desktop for locked screens
        hDesktop = OpenDesktopA("Winlogon", 0, FALSE,
            DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS);
    }

    bool success = false;
    if (hDesktop) {
        success = SetThreadDesktop(hDesktop) != FALSE;
        if (success) {
            Utils::DebugLog("SystemServiceCapture: Successfully switched to session desktop");
        }
        // Don't close desktop handle here, we'll need it
    }

    // Restore original window station
    SetProcessWindowStation(hOriginalStation);
    CloseWindowStation(hWinStation);

    return success;
}

bool SystemServiceCapture::CaptureCurrentDesktopContent(HBITMAP& hBitmap, int& width, int& height) {
    int screenWidth = GetSystemMetrics(SM_CXSCREEN);
    int screenHeight = GetSystemMetrics(SM_CYSCREEN);

    Utils::DebugLog("SystemServiceCapture: Screen dimensions: " + std::to_string(screenWidth) + "x" + std::to_string(screenHeight));

    if (screenWidth <= 0 || screenHeight <= 0) {
        m_lastError = "Invalid screen dimensions";
        return false;
    }

    // Method 1: Try GetDC(NULL) - gets DC for current desktop
    HDC hScreenDC = GetDC(NULL);
    if (!hScreenDC) {
        Utils::DebugLog("SystemServiceCapture: Failed to get screen DC");
        return false;
    }

    HDC hMemDC = CreateCompatibleDC(hScreenDC);
    if (!hMemDC) {
        m_lastError = "Failed to create memory DC";
        ReleaseDC(NULL, hScreenDC);
        return false;
    }

    HBITMAP hScreenBitmap = CreateCompatibleBitmap(hScreenDC, screenWidth, screenHeight);
    if (!hScreenBitmap) {
        m_lastError = "Failed to create bitmap";
        DeleteDC(hMemDC);
        ReleaseDC(NULL, hScreenDC);
        return false;
    }

    HGDIOBJ hOldBitmap = SelectObject(hMemDC, hScreenBitmap);

    // Force screen update first
    InvalidateRect(NULL, NULL, TRUE);
    RedrawWindow(NULL, NULL, NULL, RDW_INVALIDATE | RDW_ALLCHILDREN | RDW_UPDATENOW);

    // Try multiple BitBlt methods
    BOOL result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY);

    if (!result) {
        Utils::DebugLog("SystemServiceCapture: Standard BitBlt failed, trying CAPTUREBLT");
        result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY | CAPTUREBLT);
    }

    if (!result) {
        Utils::DebugLog("SystemServiceCapture: BitBlt with CAPTUREBLT failed, trying StretchBlt");
        result = StretchBlt(hMemDC, 0, 0, screenWidth, screenHeight,
                           hScreenDC, 0, 0, screenWidth, screenHeight, SRCCOPY);
    }

    SelectObject(hMemDC, hOldBitmap);
    DeleteDC(hMemDC);
    ReleaseDC(NULL, hScreenDC);

    if (result) {
        hBitmap = hScreenBitmap;
        width = screenWidth;
        height = screenHeight;
        Utils::DebugLog("SystemServiceCapture: Desktop content capture succeeded");
        return true;
    } else {
        DeleteObject(hScreenBitmap);
        m_lastError = "All BitBlt methods failed: " + Utils::GetLastWin32Error();
        Utils::DebugLog("SystemServiceCapture: " + m_lastError);
        return false;
    }
}